package aes

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/base64"
	"errors"
	"gitee.com/zackeus/goutil/x/encodes/secutil"
	"io"
	"sync"
)

type (
	AesCipher interface {
		Decrypt(input []byte) ([]byte, error)
		DecryptFromBase64(input string) ([]byte, error)
		Encrypt(input []byte) ([]byte, error)
		EncryptToBase64(input []byte) (string, error)
	}

	// Method 加密模式
	Method string

	Option func(c *Cipher) error

	Cipher struct {
		once      sync.Once
		key       []byte // 密钥
		iv        []byte // 初始化向量(长度必须与key相同)
		core      cipher.Block
		method    Method
		blockMode struct {
			encrypter cipher.BlockMode
			decrypter cipher.BlockMode
		}
		stream struct {
			encrypter cipher.Stream
			decrypter cipher.Stream
		}
		aead cipher.AEAD // 用于GCM模式的认证加密接口
	}
)

const (
	// MethodECB 电码本模式 密文被分割成分组长度相等的块（不足补齐）然后单独一个个加密，一个个输出组成密文
	MethodECB Method = "ECB"
	// MethodCBC 密码分组链接模式 前一个分组的密文和当前分组的明文异或或操作后再加密
	MethodCBC Method = "CBC"
	// MethodCFB 密码反馈模式
	MethodCFB Method = "CFB"
	// MethodGCM 伽罗瓦/计数器模式（认证加密模式，带标签验证）
	MethodGCM Method = "GCM"

	// GCM推荐的nonce长度（NIST推荐96位=12字节，安全性最优）
	gcmNonceSize = 12
	// GCM默认标签长度（128位=16字节）
	gcmTagSize = 16
)

func New(key []byte, opts ...Option) (AesCipher, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	c := &Cipher{
		key:    key,
		iv:     key[:block.BlockSize()],
		core:   block,
		method: MethodCBC,
	}

	for _, opt := range opts {
		if err := opt(c); err != nil {
			return nil, err
		}
	}

	if c.method != MethodGCM {
		/* 非GCM模式验证iv长度 */
		if len(c.iv) != c.core.BlockSize() {
			return nil, errors.New("IV length must equal block size")
		}
	}

	if err := c.init(); err != nil {
		return nil, err
	}
	return c, nil
}

// 初始化
func (c *Cipher) init() error {
	var err error
	c.once.Do(func() {
		switch c.method {
		case MethodCBC:
			c.blockMode.encrypter = cipher.NewCBCEncrypter(c.core, c.iv)
			c.blockMode.decrypter = cipher.NewCBCDecrypter(c.core, c.iv)
		case MethodCFB:
			c.stream.encrypter = cipher.NewCFBEncrypter(c.core, c.iv)
			c.stream.decrypter = cipher.NewCFBDecrypter(c.core, c.iv)
		case MethodECB:
			/* ECB模式不需要IV，此处无需初始化额外结构 */
		case MethodGCM:
			/* GCM使用nonce而非IV */
			/* 初始化GCM模式（包含认证功能） */
			c.aead, err = cipher.NewGCM(c.core)
			if err != nil {
				err = errors.New("init GCM failed: " + err.Error())
				return
			}
		default:
			err = errors.New("unsupported method")
		}
	})

	return err
}

// Encrypt 加密
func (c *Cipher) Encrypt(src []byte) ([]byte, error) {
	switch c.method {
	case MethodCBC:
		/* CBC */
		/* 去除补全码 */
		blockSize := c.core.BlockSize()
		padSrc := secutil.PKCS5Padding(src, blockSize)
		encrypted := make([]byte, len(padSrc))
		c.blockMode.encrypter.CryptBlocks(encrypted, padSrc)
		return encrypted, nil
	case MethodCFB:
		/* CFB */
		encrypted := make([]byte, len(src))
		c.stream.encrypter.XORKeyStream(encrypted, src)
		return encrypted, nil
	case MethodECB:
		/* ECB */
		/* 去除补全码 */
		blockSize := c.core.BlockSize()
		padSrc := secutil.PKCS5Padding(src, blockSize)
		encrypted := make([]byte, len(padSrc))
		/* 分组分块加密 */
		for bs, be := 0, blockSize; bs < len(padSrc); bs, be = bs+blockSize, be+blockSize {
			c.core.Encrypt(encrypted[bs:be], padSrc[bs:be])
		}
		return encrypted, nil
	case MethodGCM:
		/* 生成随机nonce（GCM要求每次加密使用唯一nonce） */
		nonce := make([]byte, gcmNonceSize)
		if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
			return nil, errors.New("generate nonce failed: " + err.Error())
		}
		/* 加密并生成认证标签（结果为：密文+标签） */
		ciphertext := c.aead.Seal(nil, nonce, src, nil)
		/* 最终密文格式：nonce + 密文+标签（nonce无需保密，但必须唯一） */
		return append(nonce, ciphertext...), nil
	default:
		return nil, errors.New("unsupported method")
	}
}

// Decrypt 解密
func (c *Cipher) Decrypt(enc []byte) ([]byte, error) {
	switch c.method {
	case MethodCBC:
		/* CBC */
		srcBytes := make([]byte, len(enc))
		c.blockMode.decrypter.CryptBlocks(srcBytes, enc)
		return secutil.PKCS5UnPadding(srcBytes)
	case MethodCFB:
		/* CFB */
		srcBytes := make([]byte, len(enc))
		c.stream.decrypter.XORKeyStream(srcBytes, enc)
		return srcBytes, nil
	case MethodECB:
		/* ECB 分组分块解密 */
		srcBytes := make([]byte, len(enc))
		blockSize := c.core.BlockSize()
		for bs, be := 0, blockSize; bs < len(enc); bs, be = bs+blockSize, be+blockSize {
			c.core.Decrypt(srcBytes[bs:be], enc[bs:be])
		}
		return secutil.PKCS5UnPadding(srcBytes)
	case MethodGCM:
		/* 验证密文长度（至少包含nonce和标签） */
		if len(enc) < gcmNonceSize+gcmTagSize {
			return nil, errors.New("ciphertext too short for GCM")
		}
		/* 提取nonce和密文+标签 */
		nonce := enc[:gcmNonceSize]
		ciphertextWithTag := enc[gcmNonceSize:]
		/* 解密并验证标签（标签验证失败会返回错误） */
		return c.aead.Open(nil, nonce, ciphertextWithTag, nil)
	default:
		return nil, errors.New("unsupported method")
	}
}

// DecryptFromBase64 解密Base64加密字符
func (c *Cipher) DecryptFromBase64(input string) ([]byte, error) {
	base64Decoded, err := base64.StdEncoding.DecodeString(input)
	if err != nil {
		return nil, err
	}
	return c.Decrypt(base64Decoded)
}

// EncryptToBase64 加密为base64字符
func (c *Cipher) EncryptToBase64(input []byte) (string, error) {
	d, err := c.Encrypt(input)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(d), nil
}

// WithIV 设置初始化向量
func WithIV(iv []byte) Option {
	return func(c *Cipher) error {
		c.iv = iv
		return nil
	}
}

// WithMethod 设置模式
func WithMethod(m Method) Option {
	return func(c *Cipher) error {
		c.method = m
		return nil
	}
}
